import os
import subprocess

def run_ldd(so_path):
    try:
        return subprocess.check_output(['ldd', so_path], text=True)
    except subprocess.CalledProcessError as e:
        print(f"[ldd error] {so_path}: {e}")
        return ""

def parse_ldd_output(output):
    deps = {}
    not_found = []
    for line in output.splitlines():
        if '=>' not in line:
            continue
        parts = line.split('=>')
        name = parts[0].strip()
        right = parts[1].strip()
        if 'not found' in right:
            deps[name] = None
            not_found.append(name)
        else:
            deps[name] = right.split()[0]
    return deps, not_found

def find_missing_deps(so_path, visited=None):
    visited = visited or set()
    if so_path in visited:
        return set()
    visited.add(so_path)
    output = run_ldd(so_path)
    _, not_found = parse_ldd_output(output)
    return set(not_found)

def lib_exists_in_dir(lib_name, search_dir):
    for file in os.listdir(search_dir):
        if file.startswith(lib_name):
            return True
    return False

def set_rpath(so_path, rpath="$ORIGIN"):
    try:
        subprocess.check_call(['patchelf', '--set-rpath', rpath, so_path])
        print(f"[✓] Set rpath on: {so_path}")
    except subprocess.CalledProcessError as e:
        print(f"[✗] patchelf failed for {so_path}: {e}")

def smart_rpath_fix(target_dir):
    so_files = [os.path.join(target_dir, f)
                for f in os.listdir(target_dir)
                if '.so' in f]

    print(f"\n[🔍] Scanning {len(so_files)} .so files in {target_dir}...\n")

    for so_file in so_files:
        missing = find_missing_deps(so_file)
        if not missing:
            continue

        for lib in missing:
            if lib_exists_in_dir(lib, target_dir):
                print(f"[!] Missing dep {lib} FOUND in dir, fixing {so_file}")
                set_rpath(so_file)
                break  # 一旦有一个找到了，就处理当前 .so
    print("\n[✅] Smart rpath fix done.")

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Smart RPATH fixer for dynamic libs")
    parser.add_argument("dir", help="Target directory containing .so files")
    args = parser.parse_args()

    smart_rpath_fix(args.dir)
